from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, ForeignKey, func
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
from sqlalchemy.ext.declarative import DeclarativeMeta
from sqlalchemy.exc import SQLAlchemyError
from typing import Optional, List, Dict, Any
import os
from langchain_community.utilities import SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
from langchain.llms.base import BaseLLM

Base: DeclarativeMeta = declarative_base()

class MySQLConfig:
    def __init__(self,
                 host: str = "localhost",
                 port: int = 3306,
                 user: str = "root",
                 password: str = "",
                 database: str = "financial_data"):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database

    def get_connection_string(self) -> str:
        return f"mysql+mysqlconnector://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}"

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String(100), nullable=False)
    email = Column(String(100), unique=True, nullable=False)
    created_at = Column(DateTime, default=func.now())
    transactions = relationship('Transaction', backref='user')

class Asset(Base):
    __tablename__ = 'assets'
    id = Column(Integer, primary_key=True)
    name = Column(String(100), nullable=False)
    symbol = Column(String(10), unique=True, nullable=False)
    current_price = Column(Float, nullable=False)
    transactions = relationship('Transaction', backref='asset')

class Transaction(Base):
    __tablename__ = 'transactions'
    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    asset_id = Column(Integer, ForeignKey('assets.id'), nullable=False)
    amount = Column(Float, nullable=False)
    type = Column(String(10), nullable=False)  # 'buy' or 'sell'
    timestamp = Column(DateTime, default=func.now())

class DatabaseManager:
    def __init__(self, config: MySQLConfig):
        self.engine = create_engine(config.get_connection_string())
        self.Session = sessionmaker(bind=self.engine)
        self.session = self.Session()
        
    def execute_query(self, query: str) -> List[Dict[str, Any]]:
        try:
            result = self.session.execute(query)
            return [dict(row) for row in result.mappings()]
        except SQLAlchemyError as e:
            print(f"Database error: {e}")
            return []
        finally:
            self.session.close()
            
    def get_sql_database(self) -> SQLDatabase:
        return SQLDatabase(self.engine)
        
    def create_sql_chain(self, llm: BaseLLM) -> SQLDatabaseChain:
        db = self.get_sql_database()
        return SQLDatabaseChain.from_llm(llm, db, verbose=True)